import json
import os
import sys

import backtrader as bt
import numpy as np
import pandas as pd
import plotly.graph_objects as go


class CustomStockCommission(bt.CommInfoBase):
    """
    股票佣金计算
    """

    params = (
        ("commission", 0.02),  # 0.02% 佣金比例
        ("min_commission", 5),  # 最少佣金额度，元
    )

    def getcommission(self, size: int, price: float):
        """
        按照交易额百分比commission设置佣金，最少为min_commission
        :param size:
        :param price:
        :return:
        """
        commission = abs(size) * price * self.p.commission
        return max(commission, self.p.min_commission)


class CustomFutureCommission(bt.CommInfoBase):
    """
    期货佣金计算
    """

    params = (
        ("commission", 0.01),  # 0.01%佣金比例
        ("margin", 0.12),  # 12%的保证金比例
        ("stocklike", False),  # 不是股票回测
        ("automargin", 0.12),  # 由backtrader内部逻辑决定的必须设置一个正数
        ("commtype", bt.CommInfoBase.COMM_PERC),  # 是按照比例交佣金
    )


class MultiFactorStrategy(bt.Strategy):
    params = (
        ("unit_k", 1),  # 成交单位量，
        ("is_stock", False),  # 是否是股票回测
        ("basedir", None),  # 保存回测结果的路径
        ("top_num", 100),  # 每天选取排名靠前的30只股票
        ("max_trade_cent", 0.05),  # 每天每只股票最多成交市场成交量5%
        ("trade_pos", 0.5),  # 最多用现有资金的50%买股票
    )

    def __init__(self):
        # 定义指标
        self.buy_cnt = {}
        self.sell_cnt = {}
        self.profit = {}
        self.cost_commission = {}
        self.all_cash = []
        self.every_day_choose = []

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.datetime()
        print(f"{dt.isoformat()} {txt}")

    def stop(self):
        print(
            f"总交易次数,买入{sum(self.buy_cnt.values())}：{self.buy_cnt}，\n卖出{sum(self.sell_cnt.values())}：{self.sell_cnt}"
        )
        print(f"总收益{sum(self.profit.values())}：{self.profit}")
        print(f"总佣金{sum(self.cost_commission.values())}：{self.cost_commission}")
        df = pd.DataFrame(self.all_cash)
        df.to_csv(os.path.join(self.p.basedir, "profit.csv"))
        choosedf = pd.DataFrame(self.every_day_choose)
        choosedf.to_excel(os.path.join(self.p.basedir, "choose.xlsx"))
        fig = go.Figure()
        fig.add_trace(go.Scatter(x=df["datetime"], y=df["cash"], mode="lines"))
        fig.update_layout(
            title="收益曲线图",
            xaxis_title="日期",
            yaxis_title="价值",
            xaxis_type="category",
            xaxis_rangeslider_visible=True,
        )
        fig.write_html(os.path.join(self.p.basedir, "profit.html"))

        info = {
            "profit": self.profit,
            "cost_commission": self.cost_commission,
            "buy_cnt": self.buy_cnt,
            "sell_cnt": self.sell_cnt,
        }
        with open(os.path.join(self.p.basedir, "trade_info.json"), "w") as f:
            json.dump(info, f, indent=2)

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # broker 提交/接受了，买/卖订单则什么都不做
            return

        # 检查一个订单是否完成
        # 注意: 当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    "资产%.2f，现金%.2f，标的%.2f，已买入%s, %d, 价格: %.2f, 费用: %.2f, 佣金 %.2f"
                    % (
                        self.broker.getvalue(),
                        self.broker.getcash(),
                        self.broker.getvalue() - self.broker.getcash(),
                        order.data._name,
                        order.size,
                        order.executed.price,
                        order.executed.value,
                        order.executed.comm,
                    )
                )
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.buy_cnt[order.data._name] = (
                    self.buy_cnt.get(order.data._name, 0) + 1
                )
            elif order.issell():
                self.log(
                    "资产%.2f，现金%.2f，标的%.2f，已卖出%s, %d, 价格: %.2f, 费用: %.2f, 佣金 %.2f"
                    % (
                        self.broker.getvalue(),
                        self.broker.getcash(),
                        self.broker.getvalue() - self.broker.getcash(),
                        order.data._name,
                        order.size,
                        order.executed.price,
                        order.executed.value,
                        order.executed.comm,
                    )
                )
                self.sell_cnt[order.data._name] = (
                    self.buy_cnt.get(order.data._name, 0) + 1
                )
            # 记录当前交易数量
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f"订单取消/保证金不足/拒绝 {order.data._name}")

        # 其他状态记录为：无挂起订单
        self.order = None

    # 交易状态通知，一买一卖算交易
    def notify_trade(self, trade):

        self.log(
            "交易利润%s, 毛利润 %.2f, 净利润 %.2f"
            % (trade.data._name, trade.pnl, trade.pnlcomm)
        )
        # self.log(f"交易详情：{trade.baropen}  {trade.barclose}")
        self.profit[trade.data._name] = (
            self.profit.get(trade.data._name, 0) + trade.pnlcomm
        )
        self.cost_commission[trade.data._name] = (
            self.cost_commission.get(trade.data._name, 0) + trade.pnl - trade.pnlcomm
        )

    def next(self):
        top_num = self.p.top_num
        max_trade_cent = self.p.max_trade_cent
        trade_pos = self.p.trade_pos
        now = self.datas[0].datetime.datetime()

        target_value = {
            data: data.pred[0] for data in self.datas if not np.isnan(data.pred[0])
        }
        # 获取当天每只股票的因子值

        rank_target = sorted(target_value.items(), key=lambda x: x[1], reverse=True)
        # 对股票因子值从高到低排名

        choose_target = [x[0] for x in rank_target[:top_num]]
        # 选取排名靠前的因子

        print(f"日期：{now} 选择：{sorted([x._name for x in choose_target])}")
        self.every_day_choose.append(
            {
                "datetime": now,
                "choose_stocks": ",".join([x._name for x in choose_target]),
            }
        )
        all_worth = self.broker.getvalue()
        target_pos_mm = {
            data._name: all_worth / top_num * trade_pos // (data.close[0] * 100)
            for data in choose_target
        }
        # 根据当前资产价值计算目标仓位

        cur_pos = {
            data._name: self.getposition(data).size // self.p.unit_k
            for data in self.datas
        }
        modify_pos = {
            data: target_pos_mm.get(data._name, 0) - cur_pos.get(data._name, 0)
            for data in self.datas
        }
        # 根据当前仓位和目标仓位计算买卖股票数量

        for data, mv in modify_pos.items():
            if mv < 0:
                mv = min(int(data.volume[0] * max_trade_cent), -mv)
                self.sell(data, size=self.p.unit_k * int(mv))
        for data, mv in modify_pos.items():
            if mv > 0:
                mv = min(int(data.volume[0] * max_trade_cent), mv)
                self.buy(data, size=self.p.unit_k * int(mv))

        self.all_cash.append(
            {
                "datetime": self.datas[0].datetime.datetime(),
                "cash": self.broker.getvalue(),
            }
        )


class SSPandasData(bt.feeds.PandasData):
    lines = ("pred",)  # 要添加的线
    # 设置 line 在数据源上的列位置
    params = (("pred", -1),)


def main_job(
    result_file: str,
    code_market_list: list,
    start: str,
    end: str,
    cash: float,
    slip: float,
    is_stock: bool,
):
    """
    :param result_file: 保存的模型预测结果
    :param code_market_list: 选择回测的股票代码，为空list则回测所有
    :param start: 回测时间范围“20210101”
    :param end: 结束时间范围“20240701”
    :param cash: 回测使用的金额
    :param slip: 回测滑点，0.0001表示0.01%的滑点
    :param is_stock: 是否回测股票
    :return:
    """
    dir_base = os.path.join(os.path.dirname(result_file), f"strategy_duo_tou")
    os.makedirs(dir_base, exist_ok=True)
    df = pd.read_csv(result_file)
    if "volume" not in df.columns:
        df["volume"] = 1000
    if "openInterest" not in df.columns:
        df["openinterest"] = 0.0
    else:
        df.rename(columns={"openInterest": "openinterest"}, inplace=True)
    df.drop("date_time", inplace=True, axis=1)
    df["datetime"] = pd.to_datetime(
        df["time"] + 8 * 60 * 60 * 1000, unit="ms", utc=False
    )  # 时区转换
    df.drop("time", inplace=True, axis=1)
    df.reset_index(drop=True, inplace=True)
    if code_market_list:
        df = df[df["code_market"].isin(code_market_list)].reset_index()
    df = df[
        (df["datetime"] >= pd.to_datetime(start))
        & (df["datetime"] < pd.to_datetime(end))
    ]
    df.set_index("datetime", inplace=True)
    df.sort_index(inplace=True)

    # 创建一个Cerebro引擎
    cerebro = bt.Cerebro()

    # 添加策略
    cerebro.addstrategy(
        MultiFactorStrategy,
        unit_k=100 if is_stock else 10,
        is_stock=is_stock,
        basedir=dir_base,
    )
    dt_idx = df.index.drop_duplicates()
    stock_data = {}
    for xcode, df_group in df.groupby("code_market"):
        # 添加数据
        df_new = df_group.reindex(dt_idx).ffill().bfill()
        df_new.sort_index(inplace=True)
        stock_data[xcode] = df_new
    for xcode, df_group in stock_data.items():
        data_feed = SSPandasData(dataname=df_group.sort_index(), name=xcode)
        cerebro.adddata(data_feed)

    cerebro.broker.setcash(cash)
    cmi = CustomStockCommission() if is_stock else CustomFutureCommission()
    cerebro.broker.addcommissioninfo(cmi)
    cerebro.broker.set_slippage_perc(perc=slip)
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe_ratio")
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name="annual_return")

    start_cash = cerebro.broker.getvalue()
    results = cerebro.run()
    end_cash = cerebro.broker.getvalue()
    sharpe_ratio = results[0].analyzers.sharpe_ratio.get_analysis()
    max_drawdown = results[0].analyzers.drawdown.get_analysis()
    annual_return = results[0].analyzers.annual_return.get_analysis()

    summary = dict(
        result_file=result_file,
        start=start,
        end=end,
        cash=cash,
        start_cash=start_cash,
        end_cash=end_cash,
        sharpe_ratio=sharpe_ratio,
        max_drawdown=max_drawdown,
        annual_return=annual_return,
        cmi=str(cmi.p),
        code_market_list=code_market_list,
    )
    print(summary)
    with open(os.path.join(dir_base, "backtest.json"), "w") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    print("Backtest result saved as backtest_result.png")


if __name__ == "__main__":
    datafile, start, end = sys.argv[1:]
    main_job(
        datafile,
        [],
        start,
        end,
        10000000,
        slip=0.0002,
        is_stock=True,
    )
